# -*- coding: utf-8 -*-
"""
Brief  : 书法机器人笔画编辑器
Author : 傅帅
Date   : 2020.8.14
"""
import re
import os
import json
import rospy
import tkFont
import Tkinter as tk
import tkMessageBox
import tkFileDialog
from misc.config import *
from math import *
from PIL import Image,ImageTk,ImageDraw
from geometry_msgs.msg import PoseStamped

g_root = None  # 正经的root用的是self.root,本全局变量仅用于防止重复开启
CANVAS_SIZE = 400.0

class RecordCreator():
    def __init__(self, root):
        global g_root
        if g_root != None: 
            root.destroy()
            tkMessageBox.showinfo("已启动","请勿重复开启动捕记录创建!")
            return
        g_root = root
        self.root = root
        self.is_recording = False
        self.pre_img = None         # 用以存放手动导入的预览图Image形式(已resize过)
        self.pre_imgTk = None       # 存放手动导入的预览图ImageTk形式,防止被回收
        self.record_org_list = []   # 存放动捕记录的原始数据
        self.record_list = []       # 存放xyz处理过的动捕数据,即提外接矩形+opencv坐标系(用于画图和保存)
        self.org_z = None           # 记录起始点的z作为基准
        self.min_x = float('inf')   # 记录原始数据的边界点坐标 
        self.min_y = float('inf')
        self.max_x = float('-inf')
        self.max_y = float('-inf')
        self.rect_size = [0, 0]
        self.black12 = tkFont.Font(root, family="黑体", size=12)
        self.black14 = tkFont.Font(root, family="黑体", size=14)
        self.black16 = tkFont.Font(root, family="黑体", size=16)
        self.loaded_record_name = ''
        self.update_canvas_counts = 0
        self.update_canvas_internal = UPDATE_CANVAS_INTERNAL
        self.window_dict = {}
        self.GUI_Init()
        rospy.Subscriber('vrpn_client_node/SapphiArt/pose', PoseStamped, self.nokovDataCb)
        self.root.protocol("WM_DELETE_WINDOW", self.OnClose)  # 注册关闭事件
    
    def GUI_Init(self):
        self.root.geometry('800x450')
        self.root.resizable(width = False, height = False)
        self.root.title("动捕记录创建")
        # ---------------------------------- 菜单 ----------------------------------------
        topmenu = tk.Menu(self.root, font=self.black16)
        topmenu.add_command(label="保存", command=self.SaveBtnCb)
        topmenu.add_command(label="导入", command=self.LoadRecord)
        topmenu.add_command(label="导出为笔画", command=self.ExportAsStroke)
        topmenu.add_command(label="调节Z轴深浅", command=self.AdjustZBtnCb)
        topmenu.add_command(label="等比缩放", command=self.ScaleRecordBtnCb)
        # 将topmenu设置为顶部菜单
        self.root.config(menu=topmenu)
        # ---------------------------------- 画布 ----------------------------------------
        self.canvas = tk.Canvas(self.root, bg = "#FFFFFF", width = CANVAS_SIZE, height = CANVAS_SIZE)
        self.canvas.place(x=10, y=30)
        # ---------------------------------- 按钮 ----------------------------------------
        tk.Button(self.root, text="开始记录", command=self.StartRecord, font=self.black16).place(x=480, y=90)
        tk.Button(self.root, text="结束记录", command=self.EndRecord, font=self.black16).place(x=630, y=90)
        tk.Button(self.root, text="重置记录", command=self.ResetRecord, font=self.black16).place(x=480, y=170)
        tk.Button(self.root, text="更换预览图", command=self.ChangePreImg, font=self.black16).place(x=480, y=250)
        # ---------------------------------- 标签 ----------------------------------------
        # 静态标签
        tk.Label(self.root, text="预览区:", font=self.black12, fg='black').place(x=10, y=1)
        tk.Label(self.root, text="当前动捕六轴坐标:", font=self.black12, fg='black').place(x=440, y=330)
        # 动态标签
        self.recording_state = tk.StringVar(self.root, "未在记录")
        self.recording_state_label = tk.Label(self.root, textvariable=self.recording_state, font=self.black12, fg='red')
        self.recording_state_label.place(x=344, y=1)
        self.nokov_xyz = tk.StringVar(self.root, "暂时未接收到动捕数据")
        self.nokov_rpy = tk.StringVar(self.root, "")
        tk.Label(self.root, textvariable=self.nokov_xyz, font=self.black12, fg='black').place(x=460, y=360)
        tk.Label(self.root, textvariable=self.nokov_rpy, font=self.black12, fg='black').place(x=460, y=390)
        self.size = tk.StringVar(self.root, ('尺寸:    X: {:<24.2f}Y: {:<24.2f}').format(0.0, 0.0))
        self.len = tk.StringVar(self.root, '记录点数目: 0')
        tk.Label(self.root, textvariable=self.size, font=self.black12, fg='black').place(x=450, y=10)
        tk.Label(self.root, textvariable=self.len, font=self.black12, fg='black').place(x=450, y=40)

    def OnClose(self):
        global g_root
        g_root = None
        self.root.destroy()
    
    def nokovDataCb(self, msg):
        xyz = [msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]
        xyzw = [msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z, msg.pose.orientation.w] 
        euler = list(self._quaternion2euler(xyzw))
        xyzrpy = xyz+euler
        self.nokov_xyz.set('x:{:<12.2f}y:{:<12.2f}z:{:<12.2f}'.format(xyz[0],xyz[1],xyz[2]))
        self.nokov_rpy.set('r:{:<12.2f}p:{:<12.2f}y:{:<12.2f}'.format(euler[0]*180.0/pi,euler[1]*180.0/pi,euler[2]*180.0/pi))
        if self.is_recording:  # 正在记录时，处理数据后加入列表，并更新画布
            # 使用开始记录时的第一帧更新self.org_z
            if(len(self.record_list)==0): self.org_z = xyzrpy[2]
            # 更新外接矩形边界值
            if xyzrpy[0] > self.max_x: self.max_x = xyzrpy[0]
            if xyzrpy[1] < self.min_y: self.min_y = xyzrpy[1]
            # 最小x和最大y的变化会影响record_list,需进行更新
            update_record_list = False
            if xyzrpy[0] < self.min_x: self.min_x = xyzrpy[0]; update_record_list = True
            if xyzrpy[1] > self.max_y: self.max_y = xyzrpy[1]; update_record_list = True
            if update_record_list:  # 从新更新record_list
                for i in range(len(self.record_org_list)):
                    foo = self.record_org_list[i]
                    self.record_list[i] = [foo[0]-self.min_x+PREVIEW_IMG_PADDING, self.max_y-foo[1]+PREVIEW_IMG_PADDING,\
                                           foo[2]-self.org_z, foo[3], foo[4], foo[5]]
            self.rect_size = [self.max_x - self.min_x + 2*PREVIEW_IMG_PADDING, \
                              self.max_y - self.min_y + 2*PREVIEW_IMG_PADDING]
            self.record_org_list.append(xyzrpy)
            self.record_list.append([xyzrpy[0]-self.min_x+PREVIEW_IMG_PADDING, self.max_y-xyzrpy[1]+PREVIEW_IMG_PADDING,\
                                     xyzrpy[2]-self.org_z, xyzrpy[3], xyzrpy[4], xyzrpy[5]])
            self.size.set(('尺寸:    X: {:<24.2f}Y: {:<24.2f}').format(self.rect_size[0], self.rect_size[1]))
            self.len.set('记录点数目: ' + str(len(self.record_list)))
            self.update_canvas_counts += 1
            if self.update_canvas_counts == self.update_canvas_internal:
                self.update_canvas_counts = 0
                self.UpdateCanvas()
    
    def StartRecord(self):
        if self.nokov_xyz.get() == u'暂时未接收到动捕数据': tkMessageBox.showinfo("提示","未正常接收到动捕数据!", parent=self.root);return
        if self.is_recording: tkMessageBox.showinfo("提示","已在进行记录!", parent=self.root);return
        self.ResetRecord()
        self.is_recording = True
        self.recording_state.set('正在记录')
        self.recording_state_label['fg'] = 'green'

    def EndRecord(self):
        if not self.is_recording: tkMessageBox.showinfo("提示","还未开始记录!", parent=self.root);return
        self.is_recording = False
        self.recording_state.set('未在记录')
        self.recording_state_label['fg'] = 'red'

    def ResetRecord(self):
        self.is_recording = False
        self.record_org_list = []   
        self.record_list = []       
        self.org_z = None           
        self.min_x = float('inf')   
        self.min_y = float('inf')
        self.max_x = float('-inf')
        self.max_y = float('-inf')
        self.pre_img = None
        self.pre_imgTk = None
        self.rect_size = [0, 0]
        self.nokov_xyz.set("暂时未接收到动捕数据")
        self.nokov_rpy.set("")
        self.recording_state.set('未在记录')
        self.recording_state_label['fg'] = 'red'
        self.loaded_record_name = ''
        self.size.set(('尺寸:    X: {:<24.2f}Y: {:<24.2f}').format(self.rect_size[0], self.rect_size[1]))
        self.len.set('记录点数目: ' + str(len(self.record_list)))
        self.UpdateCanvas()

    def ChangePreImg(self):
        '''手动更换预览图'''
        if len(self.record_list) == 0: tkMessageBox.showinfo("错误","工作区为空", parent=self.root);return
        if self.is_recording: tkMessageBox.showinfo("错误","正在记录中！", parent=self.root);return
        importPath = tkFileDialog.askopenfilename(title = '更换预览图', filetypes = [('笔画图片', '*.png')], parent=self.root)
        self._UpdateCanvasWithImg(img_path)

    def SaveBtnCb(self):
        if len(self.record_list) == 0: tkMessageBox.showinfo("不可保存","工作区为空", parent=self.root);return
        if self.is_recording: tkMessageBox.showinfo("不可保存","正在记录中！", parent=self.root);return
        window = tk.Toplevel(self.root)
        self.window_dict['save'] = window
        window.bind("<Return>", self.OnSaveWindowEnter)
        window.bind("<Escape>", self._OnGeneralEsc)
        window.geometry("270x60+{}+{}".format(self.root.winfo_x()+250, self.root.winfo_y()+150))
        window.title("请输入记录名称")
        self.name_entry = tk.Entry(window, width='16', font=self.black12)
        self.name_entry.place(x=15, y=20)
        self.name_entry.insert(0, self.loaded_record_name)
        submit_btn = tk.Button(window, text='保存记录', command=self.SaveRecord, font=self.black12)
        submit_btn.place(x=170, y=15)

    def OnSaveWindowEnter(self, event):
        self.SaveRecord()
    
    def SaveRecord(self):
        self._SaveRecord('record')

    def ExportAsStroke(self):
        if len(self.record_list) == 0: tkMessageBox.showinfo("不可导出","工作区为空", parent=self.root);return
        if self.is_recording: tkMessageBox.showinfo("不可导出","正在记录中！", parent=self.root);return
        window = tk.Toplevel(self.root)
        self.window_dict['export'] = window
        window.bind("<Return>", self.OnExportWindowEnter)
        window.bind("<Escape>", self._OnGeneralEsc)
        window.geometry("270x60+{}+{}".format(self.root.winfo_x()+250, self.root.winfo_y()+150))
        window.title("导出笔画名称")
        self.name_entry = tk.Entry(window, width='16', font=self.black12)
        self.name_entry.place(x=15, y=20)
        self.name_entry.insert(0, self.loaded_record_name)
        submit_btn = tk.Button(window, text='导出笔画', command=self.ExportRecord, font=self.black12)
        submit_btn.place(x=170, y=15)

    def OnExportWindowEnter(self, event):
        self.ExportRecord()
    
    def ExportRecord(self):
        self._SaveRecord('stroke')

    def LoadRecord(self):
        """导入记录或记录导出的笔画(带六轴的笔画)"""
        if self.is_recording: tkMessageBox.showinfo('错误', '正在记录中！', parent=self.root);return
        load_path = tkFileDialog.askopenfilename(title='导入记录或记录导出的笔画', filetypes=[('记录/笔画', '*.json')], parent=self.root)
        if load_path == '' or load_path == (): return
        with open(load_path, 'r') as (f):
            obj = json.load(f)
        _type = 'record' if 'record_name' in obj else 'stroke'
        self.loaded_record_name = obj[_type + '_name']
        self.rect_size = obj[_type + '_size']
        self.record_list = obj[_type + '_points']
        self.size.set(('尺寸:    X: {:<24.2f}Y: {:<24.2f}').format(self.rect_size[0], self.rect_size[1]))
        self.len.set('记录点数目: ' + str(len(self.record_list)))
        self.UpdateCanvas()

    def AdjustZBtnCb(self):
        if len(self.record_list) == 0: tkMessageBox.showinfo('错误', '工作区为空', parent=self.root);return
        if self.is_recording: tkMessageBox.showinfo('错误', '正在记录中！', parent=self.root);return
        window = tk.Toplevel(self.root)
        self.window_dict['adjust'] = window
        window.bind('<Return>', self.OnAdjustZWindowEnter)
        window.bind('<Escape>', self._OnGeneralEsc)
        window.geometry("280x60+{}+{}".format(self.root.winfo_x()+250, self.root.winfo_y()+150))
        window.title('Z调整值,负为深,正为浅 [mm]')
        self.adjustZ_entry = tk.Entry(window, width='16', font=self.black12)
        self.adjustZ_entry.place(x=15, y=20)
        self.adjustZ_entry.insert(0, '-0.0')
        submit_btn = tk.Button(window, text='调整深度', command=self.AdjustZ, font=self.black12)
        submit_btn.place(x=170, y=15)

    def OnAdjustZWindowEnter(self, event):
        self.AdjustZ()

    def AdjustZ(self):
        deltaZ = float(self.adjustZ_entry.get())
        for record in self.record_list:
            record[2] += deltaZ
        self.UpdateCanvas()
        tkMessageBox.showinfo('提示', '调整Z轴深度成功!', parent=self.root)
        self.window_dict['adjust'].destroy()

    def ScaleRecordBtnCb(self):
        if len(self.record_list) == 0: tkMessageBox.showinfo('错误', '工作区为空', parent=self.root);return
        if self.is_recording: tkMessageBox.showinfo('错误', '正在记录中！', parent=self.root);return
        window = tk.Toplevel(self.root)
        self.window_dict['scale'] = window
        window.bind('<Return>', self.OnScaleRecordWindowEnter)
        window.bind('<Escape>', self._OnGeneralEsc)
        window.geometry("270x60+{}+{}".format(self.root.winfo_x()+250, self.root.winfo_y()+150))
        window.title('请输入缩放比例[百分制]')
        self.scale_entry = tk.Entry(window, width='16', font=self.black12)
        self.scale_entry.place(x=15, y=20)
        self.scale_entry.insert(0, '100.0')
        submit_btn = tk.Button(window, text='缩放记录', command=self.ScaleRecord, font=self.black12)
        submit_btn.place(x=170, y=15)

    def OnScaleRecordWindowEnter(self, event):
        self.ScaleRecord()

    def ScaleRecord(self):
        scale_str = self.scale_entry.get()
        if re.match('^[0-9]+$|^[0-9]+\\.[0-9]*$', scale_str) == None:
            tkMessageBox.showinfo('错误', '请输入正确的百分制缩放倍率!', parent=self.root)
            return
        resize_scale = float(scale_str) / 100.0
        self.rect_size[0] *= resize_scale
        self.rect_size[1] *= resize_scale
        for record in self.record_list:
            record[0] *= resize_scale
            record[1] *= resize_scale
        tkMessageBox.showinfo('提示', '缩放记录成功!', parent=self.root)
        self.size.set(('尺寸:    X: {:<24.2f}Y: {:<24.2f}').format(self.rect_size[0], self.rect_size[1]))
        self.window_dict['scale'].destroy()

    def UpdateCanvas(self):
        '''根据self.record_list, self.rect_size等变量更新画布预览图'''
        self.canvas.delete('all')
        if( len(self.record_list)==0 ): return
        scale = 1.0    # 一像素对应多少mm
        offset = [0,0] # 预览图起始的偏置 [pix]
        if self.rect_size[0] > self.rect_size[1]:  # 横轴较长
            scale = self.rect_size[0]/CANVAS_SIZE
            offset[1] = int((CANVAS_SIZE - self.rect_size[1]/scale)/2.0)
        else:
            scale = self.rect_size[1]/CANVAS_SIZE
            offset[0] = int((CANVAS_SIZE - self.rect_size[0]/scale)/2.0)
        for record in self.record_list:
            # 算半径映射、更新预览图(画圆)
            if record[2] > 0: continue
            radius = int(MAX_RADIUS*(0-record[2])/RECORD_BRUSH_HEIGHT) if record[2]>-RECORD_BRUSH_HEIGHT else MAX_RADIUS
            self.canvas.create_oval(offset[0]+int(record[0]/scale)-radius, \
                                    offset[1]+int(record[1]/scale)-radius, \
                                    offset[0]+int(record[0]/scale)+radius, \
                                    offset[1]+int(record[1]/scale)+radius, fill='black')

    def _UpdateCanvasWithImg(self, img_path):
        # 预览图的读入和缩放,根据当前的rect_size进行缩放
        if img_path == "" or img_path == (): return
        img_org = Image.open(img_path)
        width = self.rect_size[0]
        height = self.rect_size[1]
        if height > width: #  Y轴较长
            resizeHeight = CANVAS_SIZE
            scale = CANVAS_SIZE / height
            resizeWidth = width * scale
        else: #  X轴较长
            resizeWidth = CANVAS_SIZE
            scale = CANVAS_SIZE / width
            resizeHeight = height * scale
        self.pre_img = img_org.resize((int(resizeWidth), int(resizeHeight)), Image.ANTIALIAS)
        self.pre_imgTk = ImageTk.PhotoImage(self.pre_img)
        self.canvas.delete('all')
        self.canvas.create_image(CANVAS_SIZE/2, CANVAS_SIZE/2, image=self.pre_imgTk)

    def _SaveRecord(self, _type):
        '''type:record/stroke'''
        name = self.name_entry.get()
        if name == '':
            tkMessageBox.showinfo("保存失败","文件名不可为空", parent=self.root)
            return
        # 存json
        json_obj = {
            _type+'_size':self.rect_size,
            _type+'_name':name,
            'preview_img':name+'.png',
            _type+'_points':self.record_list
        }
        if not os.path.exists('calligraphy_lib/'+_type+'s'): os.makedirs('calligraphy_lib/'+_type+'s')
        with open('./calligraphy_lib/'+_type+'s/'+name+'.json', 'w') as f:
            json_str = json.dumps(json_obj, indent=4, ensure_ascii=False)
            f.write(json_str)
        # 存图片,若有手动导入的图,则存之,否则采用自动生成的图
        if self.pre_img != None:
            # 手动导入:若self.pre_img存在,将其使用PIL保存
            self.pre_img.save('./calligraphy_lib/'+_type+'s/'+name+'.png')
        else:
            # 自动生成:若手动导入预览图不存在,直接使用PIL画图
            scale = 1.0    # 一像素对应多少mm
            if self.rect_size[0] > self.rect_size[1]:  # 横轴较长
                scale = self.rect_size[0]/CANVAS_SIZE
            else:
                scale = self.rect_size[1]/CANVAS_SIZE
            target_pic = Image.new('RGBA', [int(self.rect_size[0]/scale), int(self.rect_size[1]/scale)])
            draw = ImageDraw.Draw(target_pic)
            for record in self.record_list:
                # 算半径映射、更新预览图(画圆)
                if record[2] > 0: continue
                radius = int(MAX_RADIUS*(0-record[2])/RECORD_BRUSH_HEIGHT) if record[2]>-RECORD_BRUSH_HEIGHT else MAX_RADIUS
                draw.ellipse((int(record[0]/scale)-radius, \
                             int(record[1]/scale)-radius, \
                             int(record[0]/scale)+radius, \
                             int(record[1]/scale)+radius), fill=(0,0,0))
            target_pic.save('./calligraphy_lib/'+_type+'s/'+name+'.png')
        if _type == 'record': 
            tkMessageBox.showinfo("保存成功","已成功保存记录"+name+"到记录字库!", parent=self.root)
            self.window_dict['save'].destroy()
        else: 
            tkMessageBox.showinfo("导出成功","已成功导出记录"+name+"为笔画!", parent=self.root)
            self.window_dict['export'].destroy()

    def _OnGeneralEsc(self, event):
        event.widget.destroy()

    def _quaternion2euler(self, xyzw):
        '''四元数转欧拉角'''
        x,y,z,w = xyzw
        row = atan2(2*(w*x+y*z), 1-2*(x**2+y**2))
        pitch = asin(2*(w*y-z*x))
        yaw = atan2(2*(w*z+x*y), 1-2*(y**2+z**2))
        return row, pitch, yaw


if __name__ == "__main__":
    foo = tk.Tk()
    RecordCreator(foo)
    foo.mainloop()